package com.jfm.jfmc.plugin.snmp.utils;

import cn.hutool.core.collection.CollUtil;
import com.jfm.jfmc.plugin.snmp.entity.SNMPEntity;
import com.jfm.jfmc.plugin.snmp.entity.SNMPResponse;
import com.jfm.jfmm.collector.enums.snmp.SNMPMethod;
import com.jfm.jfmm.collector.enums.snmp.SNMPProtocol;
import lombok.extern.slf4j.Slf4j;
import org.snmp4j.*;
import org.snmp4j.event.ResponseEvent;
import org.snmp4j.mp.MPv3;
import org.snmp4j.mp.SnmpConstants;
import org.snmp4j.security.*;
import org.snmp4j.smi.*;
import org.snmp4j.transport.DefaultUdpTransportMapping;
import org.snmp4j.util.DefaultPDUFactory;
import org.snmp4j.util.TableEvent;
import org.snmp4j.util.TableUtils;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @author DYang
 * @since 2024-10-31 21:54
 *
 * SNMP 数据采集工具类
 */
@Slf4j
public class SNMPUtil {

    private static final String UDP = "udp";
    private static final String FORMAT_PATTERN =
            "{0,choice,0#|1#1 day, |1<{0,number,integer} days, }"
                    + "{1,choice,0#|1#1 hour, |1<{1,number,integer} hours, }"
                    + "{2,choice,0#|1#1 minute, |1<{2,number,integer} minutes, }"
                    + "{3,choice,0#|1#1 second, |1<{3,number,integer} seconds }";

    // 缓存实例
    private static final ConcurrentHashMap<String, Snmp> snmpCache = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<String,PDU> pduCache = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<String,Target<UdpAddress>> targetCache = new ConcurrentHashMap<>();
    private static final AtomicBoolean isListening = new AtomicBoolean(false);


    /**
     *  获取snmp数据
     *  @param entity 参数
     * @param metricId  监控项id
     */
    public static SNMPResponse getSnmp(String metricId,SNMPEntity entity) throws IOException {
        long startTime = System.currentTimeMillis();

        // 获取snmp服务
        Snmp snmp = getSnmpInstance(metricId);
        SNMPResponse snmpResponse = null;
        // 获取pdu
        PDU pdu = getPUDInstance(entity);
        // 创建目标
        Target<UdpAddress> target = getTargetInstance(metricId, entity);
        // 发送请求
        snmpResponse = send(snmp, pdu, target,entity.getMetricCode());

        // 组装数据
        snmpResponse.setOid(entity.getOids());
        snmpResponse.setHost(entity.getHost());
        snmpResponse.setCommunity(entity.getCommunity());
        snmpResponse.setStartTime(startTime);
        snmpResponse.setResponseTime(System.currentTimeMillis() - startTime);

        return snmpResponse;
    }

    /**
     *  获取snmp数据
     *  @param entity 参数
     * @param metricId  监控项id
     */
    public static SNMPResponse walkSnmp(String metricId,SNMPEntity entity) throws IOException {
        long startTime = System.currentTimeMillis();
        // 获取snmp服务
        Snmp snmp = getSnmpInstance(metricId);
        // 创建目标
        Target<UdpAddress> target = getTargetInstance(metricId, entity);
        // 发送请求
        SNMPResponse snmpResponse = walkSnmp(snmp, target, entity.getOids().split(","));

        snmpResponse.setOid(entity.getOids());
        snmpResponse.setHost(entity.getHost());
        snmpResponse.setCommunity(entity.getCommunity());
        snmpResponse.setStartTime(startTime);
        snmpResponse.setResponseTime(System.currentTimeMillis() - startTime);

        return snmpResponse;
    }

    /**
     *  获取snmp数据
     *  @param snmp snmp服务
     *  @param target 目标
     * @param oidArr  oid
     */
    public static SNMPResponse walkSnmp(Snmp snmp, Target<?> target,String ... oidArr) {
        SNMPResponse snmpResponse = new SNMPResponse();
        Map<String, String> result = new HashMap<>();

        TableUtils tableUtils = new TableUtils(snmp, new DefaultPDUFactory(PDU.GETBULK));
        OID[] oids = new OID[oidArr.length];
        for (int i = 0; i < oidArr.length; i++) {
            oids[i] = new OID(oidArr[i]);
        }

        List<TableEvent> tableEvents = tableUtils.getTable(target, oids, null, null);
        if (CollUtil.isEmpty(tableEvents)) {
            snmpResponse.setErrorMsg("snmp walk response empty content.");
            snmpResponse.setResult(result);
            return snmpResponse;
        }

        for (TableEvent tableEvent : tableEvents) {
            if (tableEvent == null || tableEvent.isError()) {
                continue;
            }
            VariableBinding[] varBindings = tableEvent.getColumns();
            int emptyCount = 0;
            for (VariableBinding binding : varBindings) {
                // when too many empty value field, ignore
                if (emptyCount > varBindings.length / 2) {
                    break;
                }
                if (binding == null) {
                    emptyCount++;
                    continue;
                }
                Variable variable = binding.getVariable();
                if (variable instanceof TimeTicks timeTicks) {
                    String value = timeTicks.toString(FORMAT_PATTERN);
                    result.put(binding.getOid().trim().toString(), value);

                } else {
                    result.put(binding.getOid().toString().trim(), bingdingHexValueToString(binding));
                }
            }
        }
        // debug 日志，方便后续追踪问题使用
        log.debug("snmp walk response: " + result);
        Map<String,String> newResult = transformValue(oidArr,result);

        snmpResponse.setResult(newResult);
        return snmpResponse;
    }

    /**
     * 转换数据，合并指标值
     * @param oidArr [0]:key [1]:value
     * @param result 原始值
     * @return
     */
    private static Map<String, String> transformValue(String[] oidArr, Map<String, String> result) {
        String key = oidArr[0];
        String value = oidArr[1];
        Map<String,String> newResult = new HashMap<>(result.size() / 2);
        for (Map.Entry<String, String> entryKey : result.entrySet()) {
            if (entryKey.getKey().startsWith(key)) {
                String key2 = entryKey.getKey();
                String valKey = value + key2.substring(key.length());

                // 获取值
                String keyValue = entryKey.getValue();
                String val = result.get(valKey);
                newResult.put(keyValue,val);

            }
        }
        return newResult;
    }

    private static String bingdingHexValueToString(VariableBinding binding) {
        // whether if binding is hex
        String hexString = binding.toValueString();
        if (hexString.contains(":")) {
            try {
                StringBuilder output = new StringBuilder();
                String[] hexArr = hexString.split(":");
                for (String hex : hexArr) {
                    output.append((char) Integer.parseInt(hex, 16));
                }
                return output.toString();
            } catch (Exception e) {
                return hexString;
            }
        } else {
            return hexString;
        }
    }

    /**
     * 初始化snmp服务
     * @param metricId   任务id
     */
    private static Snmp getSnmpInstance(String metricId) throws IOException {
        // 根据任务id生成key
        String key = metricId.substring(0,6);
        Snmp snmp = snmpCache.get(key);
        if (snmp == null) {
            synchronized (snmpCache) {
                snmp = snmpCache.get(key);
                if (snmp == null) {
                    // 创建传输映射
                    TransportMapping<? extends Address> transport = new DefaultUdpTransportMapping();
                    snmp = new Snmp(transport);
                    snmpCache.put(key, snmp);

                    // 启动监听
                    if (isListening.compareAndSet(false, true)) {
                        transport.listen();
                    }
                }
            }
        }
        return snmp;
    }

    /**
     * 初始化snmp服务
     * @param entity
     */
    private static PDU getPUDInstance(SNMPEntity entity) {
        return createPDU(entity.getMethod(), entity.getOids());
    }

    /**
     * 初始化Target 实例
     * @param key   任务id
     * @param entity 参数
     */
    private static Target<UdpAddress> getTargetInstance(String key,SNMPEntity entity) {
        // 根据任务id生成key
        Target<UdpAddress> target = targetCache.get(key);
        if (target == null) {
            synchronized (targetCache) {
                target = targetCache.get(key);
                if (target == null) {
                    // 创建PDU
                    String address = buildAddress(entity.getHost(),entity.getPort());
                    target = createTarget(entity,address);
                    targetCache.put(key, target);
                }
            }
        }
        return target;
    }

    /**
     * 发送请求
     * @param snmp snmp服务
     * @param pdu  PDU
     * @param target 目标
     * @return
     */
    private static SNMPResponse send(Snmp snmp,PDU pdu, Target<UdpAddress> target,String metricCode) throws IOException {
        Map<String,String> result = null;
        SNMPResponse snmpResponse = new SNMPResponse();
        ResponseEvent<UdpAddress> response = snmp.send(pdu, target);
        if (response != null) {
            PDU responsePDU = response.getResponse();
            if (responsePDU != null) {
                result = new HashMap<>();
                List<? extends VariableBinding> bindings = responsePDU.getVariableBindings();
                if (!bindings.isEmpty() && bindings.get(0) != null) {
                    VariableBinding vb = bindings.get(0);
                    result.put(metricCode, vb.getVariable().toString());
                }
            }
        }
        snmpResponse.setResult(result);
        return snmpResponse;
    }


    /**
     * 构建地址
     * @param host ip地址
     * @param port 端口，默认161
     * @return
     */
    private static String buildAddress(String host, Integer port) {
        return UDP + ":" + host + "/" + port;
    }

    /**
     * 创建PDU
     * @param method   get/walk
     * @param oid      OID
     */
    private static PDU createPDU(SNMPMethod method, String oid){
        PDU pdu = new PDU();
        pdu.add(new VariableBinding(new OID(oid))); // 系统描述的OID
        if ("get".equals(method.getMethod())) {
            pdu.setType(PDU.GET);
        }else {
            pdu.setType(PDU.GETNEXT);
        }
        return pdu;
    }


    /**
     * 创建目标
     * @return
     */
    private static Target<UdpAddress> createTarget(SNMPEntity entity, String address) {
        if (entity.getProtocol().getVersion() == SnmpConstants.version3) {
            return createUSMTarget(entity, address);
        } else {
            return createCommunityTarget(entity, address);
        }
    }

    private static Target<UdpAddress> createCommunityTarget(SNMPEntity entity,String address) {
        CommunityTarget<UdpAddress> target = new CommunityTarget<>();
        target.setCommunity(new OctetString(entity.getCommunity()));
        target.setAddress((UdpAddress) GenericAddress.parse(address));
        target.setRetries(entity.getRetry() == 0 ? 2 : entity.getRetry());
        target.setTimeout(entity.getTimeout() == 0 ? 1500 : entity.getTimeout());
        target.setVersion(entity.getProtocol().getVersion());
        return target;
    }

    /**
     * v3 协议认证
     * 创建USM认证目标
     */
    private static Target<UdpAddress> createUSMTarget(SNMPEntity entity, String address) {
        UserTarget<UdpAddress> target = new UserTarget<>();
        target.setAddress((UdpAddress) GenericAddress.parse(address));
        target.setRetries(2);
        target.setTimeout(1500);
        target.setVersion(SnmpConstants.version3);

        USM usm = new USM(SecurityProtocols.getInstance().addDefaultProtocols(), new OctetString(MPv3.createLocalEngineID()), 0);
        SecurityModels.getInstance().addSecurityModel(usm);

        UsmUser user = new UsmUser(new OctetString(entity.getV3().getUsername()), AuthSHA.ID, new OctetString("authPassphrase"), PrivDES.ID, new OctetString("privPassphrase"));
        usm.addUser(user.getSecurityName(), user);

        target.setSecurityLevel(SecurityLevel.AUTH_PRIV);
        target.setSecurityName(new OctetString(entity.getV3().getUsername()));
        target.setSecurityModel(3);

        return target;
    }


}
